{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gho2N54s0dXx"
      },
      "source": [
        "# SD Scripts Colab\n",
        "Created by [licyk](https://github.com/licyk)\n",
        "\n",
        "Jupyter Notebook 仓库：[licyk/sd-webui-all-in-one](https://github.com/licyk/sd-webui-all-in-one)\n",
        "\n",
        "\n",
        "## 简介\n",
        "一个在 [Colab](https://colab.research.google.com/) 部署 [sd-scripts](https://github.com/kohya-ss/sd-scripts) 的 Jupyter Notebook，可用于 Stable Diffusion 模型的训练。\n",
        "\n",
        "这个 Colab 脚本只是写来玩的，还有用来开发和测试管理模块的功能。如果要用这个脚本进行训练就参考 [SD Scripts Kaggle Jupyter NoteBook](https://github.com/licyk/sd-webui-all-in-one?tab=readme-ov-file#sd-scripts-kaggle-jupyter-notebook)，毕竟同款核心。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3rM4x-g4BiIa"
      },
      "outputs": [],
      "source": [
        "# @title 👇 参数配置\n",
        "# @markdown ## SD WebUI All In One 核心配置\n",
        "# @markdown - 模块下载地址\n",
        "SD_WEBUI_ALL_IN_ONE_URL = \"https://github.com/licyk/sd-webui-all-in-one/raw/main/sd_scripts_ipynb_core.py\" # @param {\"type\":\"string\",\"placeholder\":\"填写 SD Scripts Manager 核心下载地址\"}\n",
        "# @markdown - 强制下载核心模块即使已存在\n",
        "FORCE_DOWNLOAD_CORE = False # @param {\"type\":\"boolean\"}\n",
        "# SD WebUI All In One 功能初始化部分, 通常不需要修改\n",
        "# 如果需要查看完整代码实现, 可阅读: https://github.com/licyk/sd-webui-all-in-one/blob/main/sd_webui_all_in_one\n",
        "#################################################################################################################\n",
        "import os\n",
        "from pathlib import Path\n",
        "try:\n",
        "    _ = JUPYTER_ROOT_PATH  # type: ignore # noqa: F821\n",
        "except Exception:\n",
        "    JUPYTER_ROOT_PATH = os.getcwd()\n",
        "!python -c \"import sd_webui_all_in_one\" &> /dev/null && [ \"{FORCE_DOWNLOAD_CORE}\" != \"True\" ] || python -m pip install \"git+{SD_WEBUI_ALL_IN_ONE_URL}\"\n",
        "from sd_webui_all_in_one import logger, VERSION, SDScriptsManager\n",
        "logger.info(\"SD WebUI All In One 核心模块初始化完成, 版本: %s\", VERSION)\n",
        "# @markdown ---\n",
        "##############################################################################\n",
        "\n",
        "# @markdown ## 环境设置\n",
        "# @markdown - 工作路径, 通常不需要修改\n",
        "WORKSPACE = \"/content\" # @param {\"type\":\"string\",\"placeholder\":\"填写工作路径\"}\n",
        "# @markdown - 工作路径中文件夹名称, 通常不需要修改\n",
        "WORKFOLDER = \"sd-scripts\" # @param {\"type\":\"string\",\"placeholder\":\"填写工作文件夹\"}\n",
        "# @markdown - sd-scripts 仓库地址\n",
        "SD_SCRIPTS_REPO = \"https://github.com/kohya-ss/sd-scripts\" # @param {\"type\":\"string\",\"placeholder\":\"填写 sd-scripts 仓库的 Git 地址\"}\n",
        "# @markdown - sd-scripts 依赖文件名\n",
        "SD_SCRIPTS_REQUIREMENTS = \"requirements.txt\" # @param {\"type\":\"string\",\"placeholder\":\"填写 sd-scripts 依赖记录的文件名\"}\n",
        "# @markdown - PyTorch 版本\n",
        "TORCH_VER = \"torch==2.5.1+cu124 torchvision==0.20.1+cu124 torchaudio==2.5.1+cu124\" # @param {\"type\":\"string\",\"placeholder\":\"填写 PyTorch 软件包名和版本\"}\n",
        "# @markdown - xFormers 版本\n",
        "XFORMERS_VER = \"xformers==0.0.28.post3\" # @param {\"type\":\"string\",\"placeholder\":\"填写 xFormers 软件包名和版本\"}\n",
        "# @markdown - 使用 uv 加速 Python 软件包安装, 修改为 True 为启用, False 为禁用\n",
        "USE_UV = True # @param {type:\"boolean\"}\n",
        "# @markdown - PyPI 主镜像源\n",
        "PIP_INDEX_MIRROR = \"https://pypi.python.org/simple\" # @param {\"type\":\"string\",\"placeholder\":\"填写 PyPI 镜像地址\"}\n",
        "# @markdown - PyPI 扩展镜像源\n",
        "PIP_EXTRA_INDEX_MIRROR = \"https://download.pytorch.org/whl/cu124\" # @param {\"type\":\"string\",\"placeholder\":\"填写 PyPI 镜像地址\"}\n",
        "# @markdown - PyPI 额外镜像源\n",
        "PIP_FIND_LINKS_MIRROR = \"\" #@param {type:\"string\"}\n",
        "# @markdown - 用于下载 PyTorch 的镜像源\n",
        "PYTORCH_MIRROR = \"https://download.pytorch.org/whl/cu124\" # @param {\"type\":\"string\",\"placeholder\":\"填写 PyTorch 镜像地址\"}\n",
        "# PyPI 扩展镜像源\n",
        "PIP_FIND_LINKS_MIRROR = \"https://download.pytorch.org/whl/cu121/torch_stable.html\"\n",
        "HUGGINGFACE_MIRROR = \"https://hf-mirror.com\" # HuggingFace 镜像源\n",
        "GITHUB_MIRROR = [ # Github 镜像源\n",
        "    \"https://ghfast.top/https://github.com\",\n",
        "    \"https://mirror.ghproxy.com/https://github.com\",\n",
        "    \"https://ghproxy.net/https://github.com\",\n",
        "    \"https://gh.api.99988866.xyz/https://github.com\",\n",
        "    \"https://gh-proxy.com/https://github.com\",\n",
        "    \"https://ghps.cc/https://github.com\",\n",
        "    \"https://gh.idayer.com/https://github.com\",\n",
        "    \"https://ghproxy.1888866.xyz/github.com\",\n",
        "    \"https://slink.ltd/https://github.com\",\n",
        "    \"https://github.boki.moe/github.com\",\n",
        "    \"https://github.moeyy.xyz/https://github.com\",\n",
        "    \"https://gh-proxy.net/https://github.com\",\n",
        "    \"https://gh-proxy.ygxz.in/https://github.com\",\n",
        "    \"https://wget.la/https://github.com\",\n",
        "    \"https://kkgithub.com\",\n",
        "    \"https://gitclone.com/github.com\",\n",
        "]\n",
        "# @markdown - 检查可用的 GPU, 当 GPU 不可用时强制终止安装进程\n",
        "CHECK_AVALIABLE_GPU = False # @param {type:\"boolean\"}\n",
        "# @markdown - 重试下载次数\n",
        "RETRY = 3 # @param {type:\"slider\", min:1, max:128, step:1}\n",
        "# @markdown - 下载线程\n",
        "DOWNLOAD_THREAD = 16 # @param {type:\"slider\", min:1, max:128, step:1}\n",
        "# @markdown - 启用 TCMalloc 内存优化\n",
        "ENABLE_TCMALLOC = True # @param {type:\"boolean\"}\n",
        "#@markdown - 启用 CUDA Malloc 显存优化\n",
        "ENABLE_CUDA_MALLOC = True #@param {type:\"boolean\"}\n",
        "#@markdown - 更新内核\n",
        "UPDATE_CORE = True #@param {type:\"boolean\"}\n",
        "\n",
        "# @markdown ---\n",
        "##############################################################################\n",
        "\n",
        "# @markdown ## sd-scripts 版本设置\n",
        "# @markdown - sd-scripts 分支, 可切换成 main / dev 或者其它分支, 留空则不进行切换\n",
        "SD_SCRIPTS_BRANCH = \"dev\" # @param {\"type\":\"string\",\"placeholder\":\"填写 sd-scripts 分支名\"}\n",
        "# @markdown - 切换 sd-scripts 的版本到某个 Git 提交记录上, 留空则不进行切换\n",
        "SD_SCRIPTS_COMMIT = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写 sd-scripts 版本提交哈希值\"}\n",
        "\n",
        "# @markdown ---\n",
        "##############################################################################\n",
        "\n",
        "# @markdown ## 模型上传设置, 使用 HuggingFace / ModelScope 上传训练好的模型\n",
        "# @markdown HuggingFace: https://huggingface.co\n",
        "# @markdown ModelScope: https://modelscope.cn\n",
        "# @markdown - 使用 HuggingFace 保存训练好的模型\n",
        "USE_HF_TO_SAVE_MODEL = False # @param {type:\"boolean\"}\n",
        "# @markdown - 使用 ModelScope 保存训练好的模型\n",
        "USE_MS_TO_SAVE_MODEL = False # @param {type:\"boolean\"}\n",
        "\n",
        "# @markdown ## Token 配置, 用于上传 / 下载模型 (部分模型下载需要 Token 进行验证)\n",
        "# @markdown HuggingFace Token 在 Account -> Settings -> Access Tokens 中获取\n",
        "# @markdown - HuggingFace Token\n",
        "HF_TOKEN = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写 HuggingFace Token\"}\n",
        "# @markdown ModelScope Token 在 首页 -> 访问令牌 -> SDK 令牌 中获取\n",
        "# @markdown - ModelScope Token\n",
        "MS_TOKEN = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写 ModelScope SDK 令牌\"}\n",
        "\n",
        "# @markdown ## 用于上传模型的 HuggingFace 模型仓库的 ID, 当仓库不存在时则尝试新建一个\n",
        "# @markdown - HuggingFace 仓库的 ID (格式: \"用户名/仓库名\")\n",
        "HF_REPO_ID = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写 HuggingFace 仓库 ID\"}\n",
        "# @markdown - HuggingFace 仓库的种类\n",
        "HF_REPO_TYPE = \"model\" # @param [\"model\", \"dataset\", \"space\"]\n",
        "# @markdown HuggingFace 仓库类型和对应名称:</br>\n",
        "# @markdown model: 模型仓库</br>\n",
        "# @markdown dataset: 数据集仓库</br>\n",
        "# @markdown space: 在线运行空间仓库</br>\n",
        "\n",
        "# @markdown ## 用于上传模型的 ModelScope 模型仓库的 ID, 当仓库不存在时则尝试新建一个\n",
        "# @markdown - ModelScope 仓库的 ID (格式: \"用户名/仓库名\")\n",
        "MS_REPO_ID = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写 ModelScope 仓库 ID\"}\n",
        "# @markdown - ModelScope 仓库的种类\n",
        "MS_REPO_TYPE = \"model\" # @param [\"model\", \"dataset\", \"space\"]\n",
        "# @markdown ModelScope 仓库类型和对应名称:</br>\n",
        "# @markdown model: 模型仓库</br>\n",
        "# @markdown dataset: 数据集仓库</br>\n",
        "# @markdown space: 创空间仓库</br>\n",
        "\n",
        "# @markdown ## 设置自动创建仓库时仓库的可见性, 通常保持默认即可\n",
        "# @markdown - 设置新建的 HuggingFace 仓库可见性\n",
        "HF_REPO_VISIBILITY = False # @param {type:\"boolean\"}\n",
        "# @markdown - 设置新建的 ModelScope 仓库可见性\n",
        "MS_REPO_VISIBILITY = False # @param {type:\"boolean\"}\n",
        "\n",
        "# @markdown ## Git 信息设置, 可以使用默认值\n",
        "# @markdown - Git 的邮箱\n",
        "GIT_USER_EMAIL = \"username@example.com\" # @param {\"type\":\"string\",\"placeholder\":\"填写邮箱\"}\n",
        "# @markdown - Git 的用户名\n",
        "GIT_USER_NAME = \"username\" # @param {\"type\":\"string\",\"placeholder\":\"填写用户名\"}\n",
        "\n",
        "# @markdown ---\n",
        "##############################################################################\n",
        "\n",
        "# @markdown ## 训练日志设置, 可使用 TensorBoard / WandB 记录训练日志, 使用 WandB 可远程查看实时训练日志\n",
        "# @markdown 使用 WandB 需要填写 WANDB_TOKEN</br>\n",
        "# @markdown 如果 TensorBoard 和 WandB 同时使用, 可以改成 all</br>\n",
        "# @markdown - 使用的日志记录工具 (tensorboard / wandb / all)\n",
        "LOG_MODULE = \"tensorboard\" # @param [\"tensorboard\", \"wandb\", \"all\"]\n",
        "\n",
        "# @markdown ## WandB Token 设置\n",
        "# @markdown WandB Token 可在 https://wandb.ai/authorize 中获取\n",
        "# @markdown - WandB Token\n",
        "WANDB_TOKEN = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写 WandB Token\"}\n",
        "\n",
        "# @markdown ---\n",
        "##############################################################################\n",
        "\n",
        "# 路径设置, 通常保持默认即可\n",
        "# @markdown - 训练集保存的路径\n",
        "INPUT_DATASET_PATH = \"/content/dataset\" # @param {\"type\":\"string\",\"placeholder\":\"填写训练集保存路径\"}\n",
        "# @markdown - 训练时模型保存的路径\n",
        "OUTPUT_PATH = \"/content/working/model\" # @param {\"type\":\"string\",\"placeholder\":\"填写训练输出的模型保存路径\"}\n",
        "# @markdown - 模型下载到的路径\n",
        "SD_MODEL_PATH = \"/content/sd-models\" # @param {\"type\":\"string\",\"placeholder\":\"填写下载模型的路径\"}\n",
        "\n",
        "# @markdown ---\n",
        "##############################################################################\n",
        "\n",
        "# @markdown ## 训练模型设置, 在安装时将会下载选择的模型\n",
        "SD_MODEL = []\n",
        "\n",
        "# @markdown - Stable Diffusion 模型\n",
        "v1_5_pruned_emaonly = False # @param  {type:\"boolean\"}\n",
        "animefull_final_pruned = False # @param  {type:\"boolean\"}\n",
        "sd_xl_base_1_0_0_9vae = False # @param  {type:\"boolean\"}\n",
        "sd_xl_refiner_1_0_0_9vae = False # @param  {type:\"boolean\"}\n",
        "sd_xl_turbo_1_0_fp16 = False # @param  {type:\"boolean\"}\n",
        "animagine_xl_3_0_base = False # @param  {type:\"boolean\"}\n",
        "animagine_xl_3_0 = False # @param  {type:\"boolean\"}\n",
        "animagine_xl_3_1 = False # @param  {type:\"boolean\"}\n",
        "animagine_xl_4_0 = False # @param  {type:\"boolean\"}\n",
        "animagine_xl_4_0_opt = False # @param  {type:\"boolean\"}\n",
        "holodayo_xl_2_1 = False # @param  {type:\"boolean\"}\n",
        "kivotos_xl_2_0 = False # @param  {type:\"boolean\"}\n",
        "clandestine_xl_1_0 = False # @param  {type:\"boolean\"}\n",
        "UrangDiffusion_1_1 = False # @param  {type:\"boolean\"}\n",
        "RaeDiffusion_XL_v2 = False # @param  {type:\"boolean\"}\n",
        "kohaku_xl_delta_rev1 = False # @param  {type:\"boolean\"}\n",
        "kohakuXLEpsilon_rev1 = False # @param  {type:\"boolean\"}\n",
        "kohaku_xl_epsilon_rev2 = False # @param  {type:\"boolean\"}\n",
        "kohaku_xl_epsilon_rev3 = False # @param  {type:\"boolean\"}\n",
        "kohaku_xl_zeta = False # @param  {type:\"boolean\"}\n",
        "starryXLV52_v52 = False # @param  {type:\"boolean\"}\n",
        "heartOfAppleXL_v20 = False # @param  {type:\"boolean\"}\n",
        "heartOfAppleXL_v30 = False # @param  {type:\"boolean\"}\n",
        "sanaexlAnimeV10_v10 = False # @param  {type:\"boolean\"}\n",
        "sanaexlAnimeV10_v11 = False # @param  {type:\"boolean\"}\n",
        "SanaeXL_Anime_v1_2_aesthetic = False # @param  {type:\"boolean\"}\n",
        "SanaeXL_Anime_v1_3_aesthetic = False # @param  {type:\"boolean\"}\n",
        "Illustrious_XL_v0_1 = True # @param  {type:\"boolean\"}\n",
        "Illustrious_XL_v0_1_GUIDED = False # @param  {type:\"boolean\"}\n",
        "Illustrious_XL_v1_0 = False # @param  {type:\"boolean\"}\n",
        "Illustrious_XL_v1_1 = False # @param  {type:\"boolean\"}\n",
        "Illustrious_XL_v2_0_stable = False # @param  {type:\"boolean\"}\n",
        "jruTheJourneyRemains_v25XL = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_earlyAccessVersion = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_epsilonPred05Version = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_epsilonPred075 = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_epsilonPred077 = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_epsilonPred10Version = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_epsilonPred11Version = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPredTestVersion = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPred05Version = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPred06Version = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPred065SVersion = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPred075SVersion = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPred09RVersion = False # @param  {type:\"boolean\"}\n",
        "noobaiXLNAIXL_vPred10Version = False # @param  {type:\"boolean\"}\n",
        "ponyDiffusionV6XL_v6StartWithThisOne = False # @param  {type:\"boolean\"}\n",
        "pdForAnime_v20 = False # @param  {type:\"boolean\"}\n",
        "omegaPonyXLAnime_v20 = False # @param  {type:\"boolean\"}\n",
        "# @markdown - VAE 模型\n",
        "vae_ft_ema_560000_ema_pruned = False # @param  {type:\"boolean\"}\n",
        "vae_ft_mse_840000_ema_pruned = False # @param  {type:\"boolean\"}\n",
        "sdxl_fp16_fix_vae = True # @param  {type:\"boolean\"}\n",
        "\n",
        "v1_5_pruned_emaonly and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sd_1.5/v1-5-pruned-emaonly.safetensors\", 1])\n",
        "animefull_final_pruned and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sd_1.5/animefull-final-pruned.safetensors\", 1])\n",
        "sd_xl_base_1_0_0_9vae and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/sd_xl_base_1.0_0.9vae.safetensors\", 1])\n",
        "sd_xl_refiner_1_0_0_9vae and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/sd_xl_refiner_1.0_0.9vae.safetensors\", 1])\n",
        "sd_xl_turbo_1_0_fp16 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/sd_xl_turbo_1.0_fp16.safetensors\", 1])\n",
        "animagine_xl_3_0_base and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/animagine-xl-3.0-base.safetensors\", 1])\n",
        "animagine_xl_3_0 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/animagine-xl-3.0.safetensors\", 1])\n",
        "animagine_xl_3_1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/animagine-xl-3.1.safetensors\", 1])\n",
        "animagine_xl_4_0 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/animagine-xl-4.0.safetensors\", 1])\n",
        "animagine_xl_4_0_opt and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/animagine-xl-4.0-opt.safetensors\", 1])\n",
        "holodayo_xl_2_1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/holodayo-xl-2.1.safetensors\", 1])\n",
        "kivotos_xl_2_0 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/kivotos-xl-2.0.safetensors\", 1])\n",
        "clandestine_xl_1_0 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/clandestine-xl-1.0.safetensors\", 1])\n",
        "UrangDiffusion_1_1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/UrangDiffusion-1.1.safetensors\", 1])\n",
        "RaeDiffusion_XL_v2 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/RaeDiffusion-XL-v2.safetensors\", 1])\n",
        "kohaku_xl_delta_rev1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/kohaku-xl-delta-rev1.safetensors\", 1])\n",
        "kohakuXLEpsilon_rev1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/kohakuXLEpsilon_rev1.safetensors\", 1])\n",
        "kohaku_xl_epsilon_rev2 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/kohaku-xl-epsilon-rev2.safetensors\", 1])\n",
        "kohaku_xl_epsilon_rev3 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/kohaku-xl-epsilon-rev3.safetensors\", 1])\n",
        "kohaku_xl_zeta and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/kohaku-xl-zeta.safetensors\", 1])\n",
        "starryXLV52_v52 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/starryXLV52_v52.safetensors\", 1])\n",
        "heartOfAppleXL_v20 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/heartOfAppleXL_v20.safetensors\", 1])\n",
        "heartOfAppleXL_v30 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/heartOfAppleXL_v30.safetensors\", 1])\n",
        "sanaexlAnimeV10_v10 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/sanaexlAnimeV10_v10.safetensors\", 1])\n",
        "sanaexlAnimeV10_v11 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/sanaexlAnimeV10_v11.safetensors\", 1])\n",
        "SanaeXL_Anime_v1_2_aesthetic and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/SanaeXL-Anime-v1.2-aesthetic.safetensors\", 1])\n",
        "SanaeXL_Anime_v1_3_aesthetic and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/SanaeXL-Anime-v1.3-aesthetic.safetensors\", 1])\n",
        "Illustrious_XL_v0_1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/Illustrious-XL-v0.1.safetensors\", 1])\n",
        "Illustrious_XL_v0_1_GUIDED and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/Illustrious-XL-v0.1-GUIDED.safetensors\", 1])\n",
        "Illustrious_XL_v1_0 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/Illustrious-XL-v1.0.safetensors\", 1])\n",
        "Illustrious_XL_v1_1 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/Illustrious-XL-v1.1.safetensors\", 1])\n",
        "Illustrious_XL_v2_0_stable and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/Illustrious-XL-v2.0-stable.safetensors\", 1])\n",
        "jruTheJourneyRemains_v25XL and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/jruTheJourneyRemains_v25XL.safetensors\", 1])\n",
        "noobaiXLNAIXL_earlyAccessVersion and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_earlyAccessVersion.safetensors\", 1])\n",
        "noobaiXLNAIXL_epsilonPred05Version and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_epsilonPred05Version.safetensors\", 1])\n",
        "noobaiXLNAIXL_epsilonPred075 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_epsilonPred075.safetensors\", 1])\n",
        "noobaiXLNAIXL_epsilonPred077 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_epsilonPred077.safetensors\", 1])\n",
        "noobaiXLNAIXL_epsilonPred10Version and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_epsilonPred10Version.safetensors\", 1])\n",
        "noobaiXLNAIXL_epsilonPred11Version and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_epsilonPred11Version.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPredTestVersion and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPredTestVersion.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPred05Version and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPred05Version.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPred06Version and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPred06Version.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPred065SVersion and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPred065SVersion.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPred075SVersion and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPred075SVersion.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPred09RVersion and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPred09RVersion.safetensors\", 1])\n",
        "noobaiXLNAIXL_vPred10Version and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/noobaiXLNAIXL_vPred10Version.safetensors\", 1])\n",
        "ponyDiffusionV6XL_v6StartWithThisOne and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/ponyDiffusionV6XL_v6StartWithThisOne.safetensors\", 1])\n",
        "pdForAnime_v20 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/pdForAnime_v20.safetensors\", 1])\n",
        "omegaPonyXLAnime_v20 and SD_MODEL.append([\"https://huggingface.co/licyk/sd-model/resolve/main/sdxl_1.0/omegaPonyXLAnime_v20.safetensors\", 1])\n",
        "vae_ft_ema_560000_ema_pruned and SD_MODEL.append([\"https://huggingface.co/licyk/sd-vae/resolve/main/sd_1.5/vae-ft-ema-560000-ema-pruned.safetensors\", 1])\n",
        "vae_ft_mse_840000_ema_pruned and SD_MODEL.append([\"https://huggingface.co/licyk/sd-vae/resolve/main/sd_1.5/vae-ft-mse-840000-ema-pruned.safetensors\", 1])\n",
        "sdxl_fp16_fix_vae and SD_MODEL.append([\"https://huggingface.co/licyk/sd-vae/resolve/main/sdxl_1.0/sdxl_fp16_fix_vae.safetensors\", 1])\n",
        "\n",
        "##############################################################################\n",
        "# 下面为初始化参数部分, 不需要修改\n",
        "INSTALL_PARAMS = {\n",
        "    \"torch_ver\":TORCH_VER or None,\n",
        "    \"xformers_ver\": XFORMERS_VER or None,\n",
        "    \"git_branch\": SD_SCRIPTS_BRANCH or None,\n",
        "    \"git_commit\": SD_SCRIPTS_COMMIT or None,\n",
        "    \"model_path\": SD_MODEL_PATH or None,\n",
        "    \"model_list\": SD_MODEL,\n",
        "    \"use_uv\": USE_UV,\n",
        "    \"pypi_index_mirror\": PIP_INDEX_MIRROR or None,\n",
        "    \"pypi_extra_index_mirror\": PIP_EXTRA_INDEX_MIRROR or None,\n",
        "    \"pypi_find_links_mirror\": PIP_FIND_LINKS_MIRROR or None,\n",
        "    # Kaggle 的环境暂不需要以下镜像源\n",
        "    # \"github_mirror\": GITHUB_MIRROR or None,\n",
        "    # \"huggingface_mirror\": HUGGINGFACE_MIRROR or None,\n",
        "    \"pytorch_mirror\": PYTORCH_MIRROR or None,\n",
        "    \"sd_scripts_repo\": SD_SCRIPTS_REPO or None,\n",
        "    \"sd_scripts_requirements\": SD_SCRIPTS_REQUIREMENTS or None,\n",
        "    \"retry\": RETRY,\n",
        "    \"huggingface_token\": HF_TOKEN or None,\n",
        "    \"modelscope_token\": MS_TOKEN or None,\n",
        "    \"wandb_token\": WANDB_TOKEN or None,\n",
        "    \"git_username\": GIT_USER_NAME or None,\n",
        "    \"git_email\": GIT_USER_EMAIL or None,\n",
        "    \"check_avaliable_gpu\": CHECK_AVALIABLE_GPU,\n",
        "    \"enable_tcmalloc\": ENABLE_TCMALLOC,\n",
        "    \"enable_cuda_malloc\": ENABLE_CUDA_MALLOC,\n",
        "    \"custom_sys_pkg_cmd\": None,\n",
        "    \"update_core\": UPDATE_CORE,\n",
        "}\n",
        "HF_REPO_UPLOADER_PARAMS = {\n",
        "    \"api_type\": \"huggingface\",\n",
        "    \"repo_id\": HF_REPO_ID,\n",
        "    \"repo_type\": HF_REPO_TYPE,\n",
        "    \"visibility\": HF_REPO_VISIBILITY,\n",
        "    \"upload_path\": OUTPUT_PATH,\n",
        "    \"retry\": RETRY,\n",
        "}\n",
        "MS_REPO_UPLOADER_PARAMS = {\n",
        "    \"api_type\": \"modelscope\",\n",
        "    \"repo_id\": MS_REPO_ID,\n",
        "    \"repo_type\": MS_REPO_TYPE,\n",
        "    \"visibility\": MS_REPO_VISIBILITY,\n",
        "    \"upload_path\": OUTPUT_PATH,\n",
        "    \"retry\": RETRY,\n",
        "}\n",
        "os.makedirs(WORKSPACE, exist_ok=True)\n",
        "os.makedirs(OUTPUT_PATH, exist_ok=True)\n",
        "os.makedirs(SD_MODEL_PATH, exist_ok=True)\n",
        "os.makedirs(INPUT_DATASET_PATH, exist_ok=True)\n",
        "SD_SCRIPTS_PATH = os.path.join(WORKSPACE, WORKFOLDER)\n",
        "logger.info(\"参数设置完成\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UFVKyZcklP_G"
      },
      "outputs": [],
      "source": [
        "# @title 👇 安装环境\n",
        "logger.info(\"开始安装 sd-scripts\")\n",
        "sd_scripts = SDScriptsManager(WORKSPACE, WORKFOLDER)\n",
        "sd_scripts.install(**INSTALL_PARAMS)\n",
        "logger.info(\"sd-scripts 安装完成\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1F2_0hx2m5g8"
      },
      "outputs": [],
      "source": [
        "# @title 👇 模型下载工具: `sd_scripts.get_model()`\n",
        "# @markdown - 模型下载链接\n",
        "url = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写模型下载链接\"}\n",
        "# @markdown - 保存的文件名 (可选)\n",
        "filename = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写模型的名称\"}\n",
        "sd_scripts.get_model(\n",
        "    url=url,\n",
        "    path=SD_MODEL_PATH,\n",
        "    filename=filename if filename else None,\n",
        "    retry=RETRY,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1cMde4KuoCCJ"
      },
      "outputs": [],
      "source": [
        "# @title 👇 模型 / 训练集下载工具: `sd_scripts.repo.download_files_from_repo()`\n",
        "# @markdown 可从 HuggingFace / ModelScope 仓库下载文件\n",
        "# @markdown - 仓库类型\n",
        "api_type = \"huggingface\" # @param [\"huggingface\", \"modelscope\"]\n",
        "# @markdown - 仓库 ID\n",
        "repo_id = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写仓库的 ID\"}\n",
        "# @markdown - 仓库类型\n",
        "repo_type = \"model\" # @param [\"model\", \"dataset\", \"space\"]\n",
        "# @markdown 文件在仓库中的路径, 不填写则下载整个仓库\n",
        "folder = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写文件在仓库中的路径\"}\n",
        "sd_scripts.repo.download_files_from_repo(\n",
        "    api_type=api_type,\n",
        "    local_dir=SD_MODEL_PATH,\n",
        "    repo_id=repo_id,\n",
        "    repo_type=repo_type,\n",
        "    folder=folder,\n",
        "    retry=RETRY,\n",
        "    num_threads=DOWNLOAD_THREAD,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "clCiubeFpvO3"
      },
      "outputs": [],
      "source": [
        "# @title 👇 压缩包下载工具并解压: `sd_scripts.download_archive_and_unpack()`\n",
        "# @markdown 支持的压缩包格式为 `ZIP`, `7Z`, `TAR`\n",
        "# @markdown - 压缩包的下载链接\n",
        "url = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写压缩包的下载链接\"}\n",
        "# @markdown - 将压缩包进行重命名的名称\n",
        "name = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写压缩包进行重命名的名称\"}\n",
        "\n",
        "sd_scripts.download_archive_and_unpack(\n",
        "    url=url,\n",
        "    local_dir=INPUT_DATASET_PATH,\n",
        "    name=name if name else None,\n",
        "    retry=RETRY,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Fw3iiY3wraCf"
      },
      "outputs": [],
      "source": [
        "# @title 👇 模型 / 训练集制作工具: `make_dataset()`\n",
        "# @markdown 基于`sd_scripts.repo.download_files_from_repo()`进行封装</br>\n",
        "# @markdown 可以自动为下载好的训练集添加重复次数\n",
        "def make_dataset(\n",
        "    api_type: str,\n",
        "    local_dir: str | Path,\n",
        "    repo_id: str,\n",
        "    repo_type: str,\n",
        "    repeat: int,\n",
        "    folder: str,\n",
        ") -> None:\n",
        "    import os\n",
        "    import shutil\n",
        "    origin_dataset_path = os.path.join(local_dir, folder)\n",
        "    tmp_dataset_path = os.path.join(local_dir, f\"{repeat}_{folder}\")\n",
        "    new_dataset_path = os.path.join(origin_dataset_path, f\"{repeat}_{folder}\")\n",
        "    sd_scripts.repo.download_files_from_repo(\n",
        "        api_type=api_type,\n",
        "        local_dir=local_dir,\n",
        "        repo_id=repo_id,\n",
        "        repo_type=repo_type,\n",
        "        folder=folder,\n",
        "        retry=RETRY,\n",
        "        num_threads=DOWNLOAD_THREAD,\n",
        "    )\n",
        "    if os.path.exists(origin_dataset_path):\n",
        "        logger.info(\"设置 %s 训练集的重复次数为 %s\", folder, repeat)\n",
        "        shutil.move(origin_dataset_path, tmp_dataset_path)\n",
        "        shutil.move(tmp_dataset_path, new_dataset_path)\n",
        "    else:\n",
        "        logger.error(\"从 %s 下载 %s 失败\", repo_id, folder)\n",
        "# @markdown - 仓库类型\n",
        "api_type = \"huggingface\" # @param [\"huggingface\", \"modelscope\"]\n",
        "# @markdown - 仓库 ID\n",
        "repo_id = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写仓库 ID\"}\n",
        "# @markdown - 仓库类型\n",
        "repo_type = \"model\" # @param [\"model\", \"dataset\", \"space\"]\n",
        "# @markdown - 训练集文件夹在仓库中的路径\n",
        "folder = \"\" # @param {\"type\":\"string\",\"placeholder\":\"填写训练集文件夹在仓库中的路径\"}\n",
        "# @markdown - 设置训练集的重复次数\n",
        "repeat = 1 # @param {type:\"slider\", min:1, max:128, step:1}\n",
        "\n",
        "make_dataset(\n",
        "    api_type=api_type,\n",
        "    local_dir=INPUT_DATASET_PATH,\n",
        "    repo_id=repo_id,\n",
        "    repo_type=repo_type,\n",
        "    repeat=repeat,\n",
        "    folder=folder,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tLKy9x2mtp8A"
      },
      "outputs": [],
      "source": [
        "# @title 👇 模型训练 (可自行修改训练参数)\n",
        "\n",
        "pretrained_model_name_or_path = \"noobaiXLNAIXL_vPred10Version.safetensors\" # @param {type:\"string\"}\n",
        "vae = \"sdxl_fp16_fix_vae.safetensors\" # @param {type:\"string\"}\n",
        "train_data_dir = \"Nachoneko\" # @param {type:\"string\"}\n",
        "output_name = \"Nachoneko_2\" # @param {type:\"string\"}\n",
        "output_dir = \"Nachoneko\" # @param {type:\"string\"}\n",
        "wandb_run_name = \"Nachoneko\" # @param {type:\"string\"}\n",
        "log_tracker_name = \"lora-Nachoneko\" # @param {type:\"string\"}\n",
        "resolution = \"1024,1024\" # @param {type:\"string\"}\n",
        "save_every_n_epochs = \"1\" # @param {type:\"string\"}\n",
        "max_train_epochs = \"2\" # @param {type:\"string\"}\n",
        "train_batch_size = \"6\" # @param {type:\"string\"}\n",
        "learning_rate = \"0.0001\" # @param {type:\"string\"}\n",
        "unet_lr = \"0.0001\" # @param {type:\"string\"}\n",
        "text_encoder_lr = \"0.00001\" # @param {type:\"string\"}\n",
        "lr_scheduler = \"constant_with_warmup\" # @param {type:\"string\"}\n",
        "lr_warmup_steps = \"100\" # @param {type:\"string\"}\n",
        "optimizer_type = \"Lion8bit\" # @param {type:\"string\"}\n",
        "\n",
        "!python \"{SD_SCRIPTS_PATH}/sdxl_train_network.py\" \\\n",
        "    --pretrained_model_name_or_path=\"{SD_MODEL_PATH}/{pretrained_model_name_or_path}\" \\\n",
        "    --vae=\"{SD_MODEL_PATH}/{vae}\" \\\n",
        "    --train_data_dir=\"{INPUT_DATASET_PATH}/{train_data_dir}\" \\\n",
        "    --output_name=\"{output_name}\" \\\n",
        "    --output_dir=\"{OUTPUT_PATH}/{output_dir}\" \\\n",
        "    --wandb_run_name=\"{wandb_run_name}\" \\\n",
        "    --log_tracker_name=\"{log_tracker_name}\" \\\n",
        "    --prior_loss_weight=1 \\\n",
        "    --resolution=\"{resolution}\" \\\n",
        "    --enable_bucket \\\n",
        "    --min_bucket_reso=256 \\\n",
        "    --max_bucket_reso=4096 \\\n",
        "    --bucket_reso_steps=64 \\\n",
        "    --save_model_as=\"safetensors\" \\\n",
        "    --save_precision=\"fp16\" \\\n",
        "    --save_every_n_epochs=\"{save_every_n_epochs}\" \\\n",
        "    --max_train_epochs=\"{max_train_epochs}\" \\\n",
        "    --train_batch_size=\"{train_batch_size}\" \\\n",
        "    --gradient_checkpointing \\\n",
        "    --network_train_unet_only \\\n",
        "    --learning_rate=\"{learning_rate}\" \\\n",
        "    --unet_lr=\"{unet_lr}\" \\\n",
        "    --text_encoder_lr=\"{text_encoder_lr}\" \\\n",
        "    --lr_scheduler=\"{lr_scheduler}\" \\\n",
        "    --lr_warmup_steps=\"{lr_warmup_steps}\" \\\n",
        "    --optimizer_type=\"{optimizer_type}\" \\\n",
        "    --network_module=\"lycoris.kohya\" \\\n",
        "    --network_dim=100000 \\\n",
        "    --network_alpha=100000 \\\n",
        "    --network_args \\\n",
        "        conv_dim=100000 \\\n",
        "        conv_alpha=100000 \\\n",
        "        algo=lokr \\\n",
        "        dropout=0 \\\n",
        "        factor=8 \\\n",
        "        train_norm=True \\\n",
        "        preset=\"full\" \\\n",
        "    --optimizer_args \\\n",
        "        weight_decay=0.05 \\\n",
        "        betas=\"0.9,0.95\" \\\n",
        "    --log_with=\"{LOG_MODULE}\" \\\n",
        "    --logging_dir=\"{OUTPUT_PATH}/logs\" \\\n",
        "    --caption_extension=\".txt\" \\\n",
        "    --shuffle_caption \\\n",
        "    --keep_tokens=0 \\\n",
        "    --max_token_length=225 \\\n",
        "    --seed=1337 \\\n",
        "    --mixed_precision=\"fp16\" \\\n",
        "    --xformers \\\n",
        "    --cache_latents \\\n",
        "    --cache_latents_to_disk \\\n",
        "    --persistent_data_loader_workers \\\n",
        "    --debiased_estimation_loss \\\n",
        "    --vae_batch_size=4 \\\n",
        "    --full_fp16"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tpKPeMXN0C9X"
      },
      "outputs": [],
      "source": [
        "# @title 👇 上传模型到 HuggingFace / ModelScope\n",
        "# 使用 HuggingFace 上传模型\n",
        "if USE_HF_TO_SAVE_MODEL:\n",
        "    logger.info(\"使用 HuggingFace 保存模型\")\n",
        "    sd_scripts.repo.upload_files_to_repo(**HF_REPO_UPLOADER_PARAMS)\n",
        "\n",
        "# 使用 ModelScope 上传模型\n",
        "if USE_MS_TO_SAVE_MODEL:\n",
        "    logger.info(\"使用 ModelScope 保存模型\")\n",
        "    sd_scripts.repo.upload_files_to_repo(**MS_REPO_UPLOADER_PARAMS)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
